Fix integer overflow in RKNPU implicit bias allocation - #29249
Conversation
The RKNPU converter built implicit (all-zero) bias buffers with malloc(sizeof(T) * dim), where dim is an attacker-controlled uint32_t taken from a model weight's shape, and never checked the result before memset. On 32-bit RKNPU targets (size_t is 32-bit) a crafted dim such as 0x40000400 makes sizeof(float) * dim wrap to 4096 bytes while the tensor still advertises dim elements, yielding a heap buffer overflow when the bias is consumed. On 64-bit hosts a large dim instead produces a failing malloc whose NULL return was dereferenced by the following memset. Route all four bias sites (Conv, QLinearConv, DepthwiseConv, FC) through a single AllocZeroedBias helper that computes the byte count with SafeInt<size_t> (throws on overflow) and null-checks the allocation before zeroing.
There was a problem hiding this comment.
Pull request overview
This PR hardens the RKNPU execution provider’s ONNX-to-RKNN converter against attacker-controlled shape dimensions by preventing sizeof(T) * dim wraparound and by avoiding memset on a null allocation when creating implicit (all-zero) bias buffers.
Changes:
- Add a shared
AllocZeroedBiashelper that computes allocation sizes withSafeInt<size_t>and checksmallocresults before zeroing. - Route implicit-bias allocations in Conv/QLinearConv/DepthwiseConv/FC layer builders through the helper.
- Include ORT’s
safeint.hto use the standard overflow-checked arithmetic idiom.
Address review feedback: AllocZeroedBias now reports element_size, count, and the computed byte count in both the SafeInt-overflow and malloc-failure error paths, so an oversized/malicious model or OOM is easier to diagnose.
Reject negative and out-of-range ONNX dimensions before converting them to the uint32_t shape representation used by the RKNPU converter and DDK. This prevents malicious int64 ONNX dimensions from silently truncating at model ingestion.
Address review feedback by replacing the explicit invalid_argument throw in ToRknpuDim with ORT_ENFORCE and including core/common/common.h directly.
Keep the allocation overflow/OOM diagnostic context while replacing direct std::runtime_error throws in AllocZeroedBias with ORT_THROW and ORT_ENFORCE.
Address review feedback by replacing the raw try/catch in AllocZeroedBias with ORT_TRY/ORT_CATCH so the code remains portable to ORT_NO_EXCEPTIONS builds.
void* ptr = AllocZeroedBias(sizeof(float), dim);
free_list_.push_back(ptr); // <-- can throw (vector reallocation)If push_back throws, the allocated ptr is leaked. This is a pre-existing issue not introduced by this PR, but it's a latent leak. A safer pattern would use a std::unique_ptr.
// Allocates a zero-initialized bias buffer with overflow-checked size arithmetic.
|
Address review feedback: (1) drop the ORT_TRY/ORT_CATCH wrapper since SafeInt<size_t> already throws OnnxRuntimeException on overflow; (2) replace malloc/free_list with std::make_unique<uint8_t[]> owned by std::vector<std::unique_ptr<uint8_t[]>>, so buffers are zero-initialized and freed automatically and Clear() no longer walks raw pointers.
Address review point: drop the security-descriptive wording (malicious model / heap buffer overflow) per ORT convention; describe AllocZeroedBias neutrally as overflow-checked size arithmetic.
Description
The RKNPU execution provider's ONNX converter creates implicit (all-zero) bias
buffers when a
Conv/Gemm/QLinearConvnode omits its bias input. Thebuffer size was computed as
sizeof(T) * dim(wheredimderives from a modelweight's shape) with no overflow check, and the raw allocation was tracked in a
manually-freed
void*list.This PR hardens the converter:
int64_tdimensions are validated (viaORT_ENFORCEin a newToRknpuDimhelper) before being narrowed to the RKNPUuint32_tshape representation, rejecting negative or out-of-range values.This covers all four ingestion points (
HandleInitializer,GetInputOfOnnxModel,GetShape,GetSupportedNodes).(
AddLayerConvImpl,AddLayerQLinearConvImpl,AddLayerDepthwiseConvImpl,AddLayerFC) go through a sharedAllocZeroedBiashelper that computes thebyte count with
SafeInt<size_t>(throws on overflow) and returns azero-initialized
std::make_unique<uint8_t[]>buffer.free_list_is nowstd::vector<std::unique_ptr<uint8_t[]>>, so the bias buffers are freedautomatically and
Clear()no longer walks/free()s raw pointers.Motivation and Context
A malicious ONNX model can provide dimensions that are unsafe for the RKNPU
converter's 32-bit shape representation or for byte-size allocation arithmetic:
int64_t, while the RKNPU converter/DDK usesuint32_tshape values. Silently narrowing a large or negativeint64_tvalue can produce a misleading
uint32_tdimension.uint32_t, the originalsizeof(T) * dimcould overflowsize_ton 32-bit RKNPU targets. For example,dim = 0x40000400makessizeof(float) * dimwrap to 4096 bytes while thecreated tensor still advertises
dimelements, which would corrupt the heapwhen the bias is consumed by the driver.
mallocresult tomemsetwithout a nullcheck.
The fix uses
SafeInt<size_t>(the ORT-standard idiom for memory-sizearithmetic), validates ONNX dimensions before they enter the RKNPU
uint32_tshape model, and replaces the manual
malloc/freelist with zero-initialized,RAII-owned
std::make_unique<uint8_t[]>buffers.Validation: the RKNPU EP requires the Rockchip DDK (
RKNPU_DDK_PATH) and anARM target, so it does not build on a typical x64 dev box and has no GPU-free CI
leg. The change was validated with
clang-format,git diff --check, and astandalone test against ORT's
SafeInt.hpp, confirming the guard throws on theoverflowing dimensions (
0x40000400,0xFFFFFFFF) while preserving normal andzero-dimension behavior.
Testing: No unit test is included because the RKNPU EP is not compiled in
any CI leg (it requires the proprietary Rockchip DDK and an ARM target), and the
ToRknpuDim/AllocZeroedBiashelpers have internal (file-static) linkage,so they are not reachable from the gtest suites. The overflow/validation logic
was instead exercised with a standalone test against ORT's
SafeInt.hppasnoted above.